{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 混合精确率\n",
    "混合精度是指在训练期间在模型中同时使用 16 位和 32 位浮点类型，以使其运行速度更快，占用的内存更少。通过将模型的某些部分保持在 32 位类型中以保持数值稳定性，该模型将具有较低的步进时间，并在精度等评估指标方面进行同等训练。本指南介绍如何使用实验性的 Keras 混合精度 API 来加速模型。使用这个API可以在现代 GPU上提高性能 3 倍以上，在 TPU上提高性能 60%。\n",
    "\n",
    "> 注意：Keras 的混合精度 API 还处于实验阶段，后期可能发生改变\n",
    "\n",
    "现在，大多数模型使用 float32 数据类型，它占用 32 位内存。但是，还有两种精度较低的数据类型 float16 和 bfloat16，它们各自占用 16 位内存。现代加速器可以在 16 位数据类型中更快地运行操作，因为它们有专门的硬件来运行 16 位计算，16 位数据类型可以更快地从内存中读取。\n",
    "\n",
    "NVIDIA CPU 在 float16 中运行操作的速度比在 float32 中快，TPU 在 bfloat16 中运行操作的速度比在float32 中快。因此，这些较低精度的数据类型应尽可能在这些设备上使用。但是，出于数值原因，变量和一些计算仍应在 float32 中，以便模型训练到相同的质量。Keras 混合精度 API 允许你使用 float16 或 bfloat16 与 float32 的混合，以从 float16/bfloat16 获得性能优势，从 float32 获得数值稳定性优势。\n",
    "\n",
    "> 注意：在本指南中，\"数值稳定性\" 一词是指使用较低精度的数据类型而不是较高精度的数据类型对模型质量的影响。在 float16 或 bfloat16 中，如果在其中一种类型中运行某个操作导致模型的计算精度或其他度量指标比在float32中运行该操作更差，则我们认为该操作在 \"数值上不稳定\"。\n",
    "\n",
    "## 起步\n",
    "Keras 混合精度 API 在 tensorflow2.1 中可使用："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "\n",
    "from tensorflow import keras\n",
    "from tensorflow.keras import layers\n",
    "from tensorflow.keras.mixed_precision import experimental as mixed_precision"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 硬件支持\n",
    "然混合精度将在大多数硬件上运行，但它只会加速最新的 NVI\n",
    "DIA GPU 和 Cloud TPU 上的模型。NVIDIA GPU 支持使用 float16 和 float32 的混合，而 TPU 支持 bfloat16 和 float32 的混合。\n",
    "\n",
    "在 NVIDIA GPU中，那些具有 7.0 或更高计算能力的 GPU 将从混合精度中得到最大的性能优势，因为它们有特殊的硬件单元，称为张量核，用于加速 float16 矩阵乘法和卷积。较旧的 GPU 对于使用混合精度没有数学性能优势，但是可以通过节省内存和带宽可以提高一些速度。你可以在 NVIDIA 的 [CUDA GPU网页](https://developer.nvidia.com/cuda-gpus)上查找 GPU 的计算能力。从混合精度中受益最大的 GPU 示例：包括 RTX GPU、Titan V 和 V100。\n",
    "\n",
    "> 注意：如果在 Google Colab 中运行本指南，GPU 运行时通常会连接一个 P100。P100 的计算能力为 6.0，预计不会出现显著的加速。\n",
    "\n",
    "你可以用以下方法检查你的 GPU 类型。此命令仅在安装了 NVIDIA 驱动程序时才存在，否则，以下命令将引发错误：\n",
    "```sh\n",
    "nvidia-smi -L\n",
    "```\n",
    "\n",
    "所有的 Cloud TPU 都支持 bfloat16.\n",
    "\n",
    "即使是在 CPU 和 老式 GPU上，也不需要加速，混合精度的 API 仍然可以用于单元测试、调试，或者仅仅用于测试API。\n",
    "\n",
    "## 设置 dtype 策略\n",
    "要在Keras中使用混合精度，需要创建 [tf.Keras.mixed_precision.experimental.Policy](https://tensorflow.google.cn/api_docs/python/tf/keras/mixed_precision/experimental/Policy)，通常称为 dtype Policy。数据类型策略指定将在其中运行的数据类型层。在本指南中，你将从字符串 \"mixed_float16\" 构造一个策略，并将其设置为全局策略。这将导致随后创建的层使用 float16 和 float32 的混合精度。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "policy = mixed_precision.Policy('mixed_float16')\n",
    "mixed_precision.set_policy(policy)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "该策略指定了一个层的两个重要方面：该层的计算是在dtype中完成的，以及该层的变量是在 dtype 中完成的。在上面，你创建了一个 mixed_float16 策略（）即 mixed_precision。通过将字符串 \"mixed_float16\" 传递给它的构造函数而创建的策略）。使用此策略，层将使用 float16  计算和 而使用 float32  变量。为了提高性能，计算是在  float16  中完成的，但是为了保持数值的稳定性，变量必须保存在 float32 中。你可以直接查询策略的这些属性。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Compute dtype: float16\n",
      "Variable dtype: float32\n"
     ]
    }
   ],
   "source": [
    "print('Compute dtype: %s' % policy.compute_dtype)\n",
    "print('Variable dtype: %s' % policy.variable_dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "如前所述，mixed_float16 策略将极大地提高 NVIDIA GPU 上的性能，计算能力至少为 7.0。该策略将在其他 GPU 和 CPU 上运行，但可能不会提高性能。对于 TPU，应该使用 mixed_bfloat16 策略。\n",
    "\n",
    "## 构建模型\n",
    "接下来，让我们开始构建一个简单的模型。非常小的模型通常不会从混合精度中受益，因为 tensorflow 运行时的开销通常控制了执行时间，使得 GPU 上的任何性能改进都可以忽略不计。因此，如果使用 GPU，我们将构建两个大而密集的层，每个层有 4096 个单元。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "该模型将在一个CPU上运行64个单元\n"
     ]
    }
   ],
   "source": [
    "inputs = keras.Input(shape=(784,), name='digits')\n",
    "if tf.config.list_physical_devices('GPU'):\n",
    "  print('The model will run with 4096 units on a GPU')\n",
    "  num_units = 4096\n",
    "else:\n",
    "  # 在 CPU 上使用更少的单元，这样模型可以在合理的时间内完成\n",
    "  print('该模型将在一个CPU上运行64个单元')\n",
    "  num_units = 64\n",
    "dense1 = layers.Dense(num_units, activation='relu', name='dense_1')\n",
    "x = dense1(inputs)\n",
    "dense2 = layers.Dense(num_units, activation='relu', name='dense_2')\n",
    "x = dense2(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "每个层都有一个策略，默认情况下使用全局策略。因此，每个稠密层都有 mixed_float16 策略，因为之前将全局策略设置为 mixed_float16。这将导致稠密层执行浮点计算，并拥有浮点 32 个变量。为了进行浮点计算，它们将输入转换为 float16，结果导致输出为 float16。它们的变量是 float32，在调用层时将转换为float16，以避免由于dtype不匹配而导致的错误。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x.dtype: float16\n",
      "dense1.kernel.dtype: float32\n"
     ]
    }
   ],
   "source": [
    "print('x.dtype: %s' % x.dtype.name)\n",
    "# 'kernel'是dense1的变量\n",
    "print('dense1.kernel.dtype: %s' % dense1.kernel.dtype.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，创建输出预测。通常，你可以创建如下所示的输出预测，但是对于 float16，这在数值上并不总是稳定的。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "输出 dtype: float16\n"
     ]
    }
   ],
   "source": [
    "# softmax 和模型输出将是 float16，而它应该是 float32\n",
    "outputs = layers.Dense(10, activation='softmax', name='predictions')(x)\n",
    "print('输出 dtype: %s' % outputs.dtype.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在模型末尾的 softmax 激活应该是 float32。因为 dtype 策略是 mixed_float16，所以 softmax 激活通常会有一个 float16 计算 dtype 并输出一个 float16 张量。\n",
    "这可以通过分离密集层和软 max 层以及将 dtype='float32 '传递给软 max 层来修复："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Outputs dtype: float32\n"
     ]
    }
   ],
   "source": [
    "# softmax 和模型输出将是 float16\n",
    "x = layers.Dense(10, name='dense_logits')(x)\n",
    "outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)\n",
    "print('Outputs dtype: %s' % outputs.dtype.name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "将 dtype='float32' 传递给 softmax 层构造函数将覆盖该层的 dtype 策略，使之成为 float32 策略，该策略执行计算并将变量保存在 float32 中。同样，我们也可以传递 dtype=mixed_precision.Policy('float32')，层总是将 dtype 参数转换为策略。因为激活层没有变量，所以策略的变量 dtype 被忽略，但是策略的计算数据类型 float32 会导致 softmax 和模型输出为 float32。\n",
    "\n",
    "在模型的中间添加一个 float16 softmax 是可以的，但是在模型的末尾添加的 softmax 应该是 float32。原因是，如果从 softmax 到损耗的中间张量是 float16 或 bfloat16，可能会出现数值问题。\n",
    "\n",
    "你可以通过传递 dtype='float32' 来将任何层的 dtype 覆盖为 float32，如果你认为它在进行 float16 计算时在数值上不稳定的话。但是通常，这只在模型的最后一层是必要的，因为大多数层使用 mixed_float16 和mixed_bfloat16 都有足够的精度。\n",
    "\n",
    "即使模型没有以 softmax 结束，输出仍然应该是 float32。虽然这个特定的模型不需要，但是模型输出可以通过以下方式转换为float32:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 线性激活是一个恒等函数。因此，这只是将 '输出\" 转换为 float32。\n",
    "# 在这个特殊的例子中，\"输出\" 已经是float32了，所以这是无操作的。\n",
    "outputs = layers.Activation('linear', dtype='float32')(outputs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 接下来，完成并编译模型，并生成输入数据。\n",
    "model = keras.Model(inputs=inputs, outputs=outputs)\n",
    "model.compile(loss='sparse_categorical_crossentropy',\n",
    "              optimizer=keras.optimizers.RMSprop(),\n",
    "              metrics=['accuracy'])\n",
    "\n",
    "(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\n",
    "x_train = x_train.reshape(60000, 784).astype('float32') / 255\n",
    "x_test = x_test.reshape(10000, 784).astype('float32') / 255"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本例将输入数据 从int8 转换为 float32。我们不强制转换为 float16，因为 255 的除法是在 CPU 上的，它运行的 float16 操作比 float32 慢。在这种情况下，性能差异可以忽略不计，但一般来说，如果在CPU上运行，则应该在 float32 中运行输入处理数学。模型的第一层将把输入转换为 float16，因为每一层都将浮点输入转换为计算数据类型。\n",
    "\n",
    "获取模型的初始权重。这将允许从零开始的训练再次加载重量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "initial_weights = model.get_weights()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 使用 Model.fit 训练模型\n",
    "下一步，训练模型："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Train on 48000 samples, validate on 12000 samples\n",
      "Epoch 1/5\n",
      "48000/48000 [==============================] - 58s 1ms/sample - loss: 1.9588 - accuracy: 0.4247 - val_loss: 1.4600 - val_accuracy: 0.6440\n",
      "Epoch 2/5\n",
      "48000/48000 [==============================] - 52s 1ms/sample - loss: 1.2476 - accuracy: 0.7277 - val_loss: 0.9527 - val_accuracy: 0.7817\n",
      "Epoch 3/5\n",
      "48000/48000 [==============================] - 55s 1ms/sample - loss: 0.8721 - accuracy: 0.7974 - val_loss: 0.7185 - val_accuracy: 0.8354\n",
      "Epoch 4/5\n",
      "48000/48000 [==============================] - 53s 1ms/sample - loss: 0.6981 - accuracy: 0.8284 - val_loss: 0.6075 - val_accuracy: 0.8518\n",
      "Epoch 5/5\n",
      "48000/48000 [==============================] - 55s 1ms/sample - loss: 0.5893 - accuracy: 0.8518 - val_loss: 0.5067 - val_accuracy: 0.8808\n",
      "10000/10000 - 7s - loss: 0.5183 - accuracy: 0.8745\n",
      "测试的损失值: 0.5183230932235717\n",
      "测试的精度: 0.8745\n"
     ]
    }
   ],
   "source": [
    "history = model.fit(x_train, y_train,\n",
    "                    batch_size=8192,\n",
    "                    epochs=5,\n",
    "                    validation_split=0.2)\n",
    "test_scores = model.evaluate(x_test, y_test, verbose=2)\n",
    "print('测试的损失值:', test_scores[0])\n",
    "print('测试的精度:', test_scores[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "注意：模型在日志中打印每个样本的时间：例如，\"4us/sample\"。由于 tensorflow 花了一段时间优化模型，第一个周期可能较慢，但之后每个样本的时间应该稳定。\n",
    "\n",
    "如果在 Colab 中运行此指南，则可以将混合精度的性能与 float32 进行比较。为此，要在设置dtype策略部分将策略从 mixed_float16 更改为 float32，然后重新运行所有单元格到现在。在至少具有计算能力 7.0 的GPU 上，你应该会看到每个样本的时间显著增加，这表明混合精度加快了模型的速度。例如，对于 Titan V GPU，每个采样时间从 4us 增加到 12us。确保将策略更改回 mixed_float16，并在继续使用指南之前重新运行单元格。\n",
    "\n",
    "对于许多实际模型，混合精度还允许在不耗尽内存的情况下将批处理大小加倍，因为 float16 张量占用了一半内存。但是，这不适用于这个简单模型，因为你可以在任何数据类型中运行该模型，其中每个批由 60000个图像的整个 MNIST 数据集组成。\n",
    "\n",
    "如果在 TPU 上运行混合精度，与在 GPU 上运行混合精度相比，你将看不到多少性能提高。这是因为即使使用默认的 dtype 策略 float32，TPU 也已经在底层的 bfloat16 中执行了某些操作。TPU硬件不支持 float32 用于某些在 bfloat16 中数值稳定的操作，例如 matmul。对于这种操作，TPU 后端将在内部默认地使用bfloat16。因此，将 dtype='float32' 传递给使用这种操作层可能没有数值效果，但是用 bfloat16 计算运行这种层也没什么害处。\n",
    "\n",
    "## 损失比例\n",
    "损耗定标是 tf.keras.Mode l中的一种技术。fit 使用 mixed_float16 策略自动执行，以避免数值下溢。本节将介绍损失扩展以及如何自定义其行为。\n",
    "### 下溢和上溢\n",
    "与 float32 相比，float16 数据类型的动态范围很小。这意味着上面的值将溢出到无穷大，而下面的值将下溢到零。float32 和 bfloat16的动态范围要大得多，所以溢出和下溢都不是问题。比如："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = tf.constant(256, dtype='float16')\n",
    "(x ** 2).numpy()  # 上溢\n",
    "\n",
    "x = tf.constant(1e-5, dtype='float16')\n",
    "(x ** 2).numpy()  # 下溢"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "实际上，使用 float16 溢出的情况很少发生。此外，下溢也很少发生在向前传播。然而，在反向过程中，梯度可能会下溢到零。损失定标是一种技术，以防止这种下溢。\n",
    "\n",
    "## 损失扩大\n",
    "损失比例的基本概念很简单：简单地将损失乘以某个较大的数。我们称这个数字为损失等级。这将导致梯度也按比例缩放，大大减少下溢的机会。计算完最终梯度变化后，将其除以，使其返回到正确的值。\n",
    "\n",
    "这个过程的伪代码是:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'compute_gradient' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-15-2fbb07dcf757>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      3\u001b[0m \u001b[0mloss\u001b[0m \u001b[1;33m*=\u001b[0m \u001b[0mloss_scale\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      4\u001b[0m \u001b[1;31m# 我们假设 \"梯度\" 是 float32。我们不想划分梯度为 float16\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 5\u001b[1;33m \u001b[0mgrads\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0mcompute_gradient\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0mloss\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mmodel\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mtrainable_variables\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m      6\u001b[0m \u001b[0mgrads\u001b[0m \u001b[1;33m/=\u001b[0m \u001b[0mloss_scale\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mNameError\u001b[0m: name 'compute_gradient' is not defined"
     ]
    }
   ],
   "source": [
    "loss_scale = 1024\n",
    "loss = model(inputs)\n",
    "loss *= loss_scale\n",
    "# 我们假设 \"梯度\" 是 float32。我们不想划分梯度为 float16\n",
    "grads = compute_gradient(loss, model.trainable_variables)\n",
    "grads /= loss_scale"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "选择一个损失范围是很困难的。如果损失尺度过低，梯度仍然可能下溢到零。如果过高，则会出现相反的问题：梯度可能会溢出到无穷大。\n",
    "\n",
    "为了解决这个问题，tensorflow 动态地确定损失的范围，所以你不必手动选择。如果你使用tf.keras.Model，则会为你完成损失缩放，因此你不必做任何额外的工作。这将在下一节中进一步解释。\n",
    "\n",
    "## 选择损失等级\n",
    "每个 dtype 策略都有一个关联的 [tf.mixed_precision.experimental.losscale](https://tensorflow.google.cn/api_docs/python/tf/mixed_precision/experimental/LossScale) 对象，该对象表示一个固定的或动态的损耗标度。默认情况下，mixed_float16 策略的损失比例是[tf.mixed_precision.experiment.DynamicLossScale](https://tensorflow.google.cn/api_docs/python/tf/mixed_precision/experimental/DynamicLossScale)，它动态确定损失比例值。默认情况下，其他策略没有损失比例，因为只有在使用 float16 时才有必要。你可以查询策略的损失等级："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "损失等级: DynamicLossScale(current_loss_scale=32768.0, num_good_steps=30, initial_loss_scale=32768.0, increment_period=2000, multiplier=2.0)\n"
     ]
    }
   ],
   "source": [
    "loss_scale = policy.loss_scale\n",
    "print('损失等级: %s' % loss_scale)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "损耗表打印了许多内部状态，但是你可以忽略它。最重要的部分是 current_loss_scale部分，它显示了 loss scale 的当前值。\n",
    "\n",
    "你可以在构建 dtype 策略时通过传递一个数字来使用静态损失范围。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FixedLossScale(1024.0)\n"
     ]
    }
   ],
   "source": [
    "new_policy = mixed_precision.Policy('mixed_float16', loss_scale=1024)\n",
    "print(new_policy.loss_scale)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "dtype 策略构造函数总是将丢失比例转换为 losscale 对象。在本例中，它被转换为[tf.mixed_precision.experimental.fixedLossCalle](https://tensorflow.google.cn/api_docs/python/tf/mixed_precision/experimental/FixedLossScale)，这是除了 DynamicLossScale 之外唯一的其他LossCalle 子类。\n",
    "\n",
    "> 注意：不建议使用除动态损失等级以外的任何方法。选择一个固定的损失比例可能是困难的，因为将它设置得太低也会导致模型不能进行训练，而将它设置得太高则会导致 Infs 或 NaNs 出现在梯度中。动态损失等级通常接近最优损失等级，因此你不必做任何工作。目前，动态损失表比固定损失表慢一些，但性能在未来会有所提高。\n",
    "\n",
    "模型和层一样，都有一个数据类型策略。如果存在，模型将使用其策略的损失比例在 tf.keras.model.fit 方法中应用损失比例。这意味着，如果使用 Model.fit，则完全不必担心损失缩放：默认情况下，mixed_float16策略将具有动态损失缩放，Model.fit 将应用它。\n",
    "\n",
    "使用自定义培训循环时，模型将忽略策略的损失比例，你必须手动应用它。下一节将对此进行解释。\n",
    "\n",
    "## 使用自定义的训练循环训练模型\n",
    "到目前为止，你使用tf.Keras.model.fit训练了一个混合精度的 Keras 模型。接下来，你将使用混合精度和自定义训练循环。如果还不知道自定义培训循环是什么，先阅读 [Custom training guide](https://tensorflow.google.cn/tutorials/customization/custom_training_walkthrough)。\n",
    "\n",
    "在 float32 中运行一个混合精度的自定义训练循环需要做两个改变：\n",
    " 1. 以混合精度构建模型。\n",
    "2. 如果使用 mixed_float16，则显式地使用损失缩放。\n",
    "\n",
    "对于步骤 2，你将使用 [tf.keras.mixed_precision.experimental](https://tensorflow.google.cn/api_docs/python/tf/keras/mixed_precision/experimental/LossScaleOptimizer)，包装优化器并应用损失缩放的类。它有两个参数：优化器和损失规模。按照下列方法建立一个动态损失等级："
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = keras.optimizers.RMSprop()\n",
    "optimizer = mixed_precision.LossScaleOptimizer(optimizer, loss_scale='dynamic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "传递 \"dynamic\" 等同于传递 tf.mixed_precision.experimental.DynamicLossScale()。\n",
    "\n",
    "接下来，定义 loss 对象和 [tf.data.Datasets](https://tensorflow.google.cn/api_docs/python/tf/data/Dataset)。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "loss_object = tf.keras.losses.SparseCategoricalCrossentropy()\n",
    "train_dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
    "                 .shuffle(10000).batch(8192))\n",
    "test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(8192)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来，定义训练步骤函数。从损失规模优化的两个新方法来扩大损失和取消规模的梯度：\n",
    "1. get_scaled_loss(loss)：将损失乘以损失等级。\n",
    "2. get_unscaled_gradients(gradients): 将一个按比例缩放的梯度列表作为输入，并将每个梯度除以损失比例以取消缩放。\n",
    "\n",
    "必须使用这些函数以防止梯度中的下溢。losscaleumizer.apply_gradients 将在没有 Infs 或 NaNs 的情况下应用渐变。它还将更新损失等级，如果梯度具有 Infs 或 NaNs，则将其减半，如果梯度不具有 Infs 或 NaNs，则可能增加损失规模。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tf.function\n",
    "def train_step(x, y):\n",
    "  with tf.GradientTape() as tape:\n",
    "    predictions = model(x)\n",
    "    loss = loss_object(y, predictions)\n",
    "    scaled_loss = optimizer.get_scaled_loss(loss)\n",
    "  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)\n",
    "  gradients = optimizer.get_unscaled_gradients(scaled_gradients)\n",
    "  optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
    "  return loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "LossScaleOptimizer 可能会在训练开始时跳过最初的几个步骤。损失标度一开始就很高，这样就可以很快确定最佳损失标度。几步后，损失规模将趋于稳定，很少有步骤将被跳过。此过程是自动发生的，不会影响训练质量。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "周期 0: 损失值=1.8140946626663208, 测试精度=0.7184000015258789\n",
      "周期 1: 损失值=1.0170257091522217, 测试精度=0.8263999819755554\n",
      "周期 2: 损失值=0.70814448595047, 测试精度=0.8600000143051147\n",
      "周期 3: 损失值=0.5634843111038208, 测试精度=0.8472999930381775\n",
      "周期 4: 损失值=0.4858713746070862, 测试精度=0.8849999904632568\n"
     ]
    }
   ],
   "source": [
    "# 定义训练步长\n",
    "@tf.function\n",
    "def test_step(x):\n",
    "  return model(x, training=False)\n",
    "\n",
    "# 加载模型的初始权重，这样你就可以从头开始重新训练。\n",
    "model.set_weights(initial_weights)\n",
    "\n",
    "# 最后，运行自定义循环训练\n",
    "for epoch in range(5):\n",
    "  epoch_loss_avg = tf.keras.metrics.Mean()\n",
    "  test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
    "      name='test_accuracy')\n",
    "  for x, y in train_dataset:\n",
    "    loss = train_step(x, y)\n",
    "    epoch_loss_avg(loss)\n",
    "  for x, y in test_dataset:\n",
    "    predictions = test_step(x)\n",
    "    test_accuracy.update_state(y, predictions)\n",
    "  print('周期 {}: 损失值={}, 测试精度={}'.format(epoch, epoch_loss_avg.result(), test_accuracy.result()))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GPU 性能技巧\n",
    "下面是在 GPU上使用混合精度时的一些性能提示：\n",
    "### 增加批大小\n",
    "如果不影响模型质量，可尝试在使用混合精度时以两倍于批大小的方式运行。由于 float16 张量占用一半内存，这通常允许你在不耗尽内存的情况下将批处理大小增加一倍。增加批处理大小通常会增加训练吞吐量，即模型运行时每秒的训练元素数。\n",
    "### 确保 GPU 张量核被使用\n",
    "如前所述，现代 NVIDIA GPU 使用一种称为 Tensor Cores 的特殊硬件单元，可以非常快地进行 float16 矩阵乘法。然而，张量核要求张量的某些维度是 8 的倍数。在下面的例子中，一个参数是粗体的当且仅当它需要 8 的倍数才能使用张量核时：\n",
    "1. tf.keras.layers.Dense(units=64)\n",
    "2. tf.keras.layers.Conv2d(filters=48, kernel_size=7, stride=3)\n",
    "    对于其他卷积层也是类似的，比如 tf.keras.layers.Conv3d\n",
    "3. tf.keras.layers.LSTM(units=64)\n",
    "    对于其他的 RNN 也一样，比如 tf.keras.layer.GUR\n",
    "4. tf.keras.Model.fit(epochs=2, batch_size=128)\n",
    "\n",
    "你应该尽量使用张量核。如果你想了解更多，[ NVIDIA deep learning performance guide](https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html)描述了使用张量核以及其他张量核相关性能信息的确切要求。\n",
    "### XLA\n",
    "XLA 是一个编译器，它可以进一步提高混合精度性能，也可以在较小程度上提高 float32 性能。详情请参阅[XLA指南](https://tensorflow.google.cn/xla)。\n",
    "\n",
    "## Cloud TPU 性能技巧\n",
    "与在 GPU 上一样，你应该尝试将批处理大小加倍，因为 bfloat16 张量使用了一半的内存。批量大小翻倍可能会增加培训吞吐量。\n",
    "\n",
    "\n",
    "TPU 不需要任何其他混合的特定于精度的调优来获得最佳性能。TPU 已经要求使用 XLA。它们受益于 128 的倍数，但是这同样适用于 float32，就像它适用于混合精度一样。有关一般 TPU 性能技巧，请参阅 [Cloud TPU Performance Guide](https://cloud.google.com/tpu/docs/performance-guide)，这些技巧适用于混合精度和 float32。\n",
    "\n",
    "## 总结\n",
    "1. 如果使用至少具有计算能力 7.0 的 TPU 或 NVIDIA GPU，则应该使用混合精度，因为它可以将性能提高 3倍。\n",
    "2. 可以对以下使用混合精度：\n",
    "```python\n",
    "policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16')\n",
    "mixed_precision.set_policy(policy)\n",
    "```\n",
    "3. 如果你的模型以 softmax 结束，确保它的维度是 float32。无论你的模型以什么结尾，都要确保输出是float32。\n",
    "4. 如果使用带 mixed_float16 的自定义训练循环，除了上述行之外，还需要使用tf.keras.mixed_precision.experimental.losscaleumizer 包装优化器。然后调用optimizer.get_scaled_loss 来缩放损失，并调用 optimizer.get_unscaled_gradients 来取消梯度缩放。\n",
    "5. 如果不降低评估精度，则将训练批大小加倍。\n",
    "6. 在 GPU 上，确保大多数张量维是 8 的倍数，以最大化性能。\n",
    "\n",
    "有关使用 tf.keras.mixed_precision AP I的混合精度的更多示例，请参见[官方模型库](https://github.com/tensorflow/models/tree/master/official)。大多数官方模型，如 ResNet 和Transformer，都将通过传递——dtype=fp16，使用混合精度运行。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
